// Copyright (c) 2019 Presto Labs Pte. Ltd.
// Author: jaewon

#include "coin2/base/math/time_moving_window.h"

#include <string>
#include <vector>

#include <gtest/gtest.h>
#include <xtensor/xarray.hpp>

namespace coin2::base::math {

TEST(TimeMovingWindowTest, Basic) {
  SimpleTimeMovingWindow<int> window;

  window.Push(0, 0);
  EXPECT_EQ(window.size(), 0);
  EXPECT_TRUE(window.empty());
  EXPECT_EQ(window.num_popped(), 1);
}

TEST(TimeMovingWindowTest, Basic2) {
  SimpleTimeMovingWindow<int> window;
  window.ResizeWindow(1);
  EXPECT_EQ(window.window_size(), 1);

  window.Push(0, 1);
  window.Push(0, 2);
  window.Push(0, 3);
  window.Push(0, 4);
  EXPECT_EQ(window.size(), 4);
  EXPECT_FALSE(window.empty());
  EXPECT_EQ(window.num_popped(), 0);

  window.Push(5, 5);
  EXPECT_EQ(window.size(), 1);
  EXPECT_EQ(window.num_popped(), 4);

  window.PushTime(5);
  EXPECT_EQ(window.size(), 1);
  EXPECT_EQ(window.num_popped(), 4);

  window.PushTime(6);
  EXPECT_EQ(window.size(), 0);
  EXPECT_EQ(window.num_popped(), 5);
}

TEST(TimeMovingWindowTest, LastValueWrapperTest) {
  SimpleTimeMovingWindow<int, true> window;
  window.ResizeWindow(1);
  EXPECT_EQ(window.window_size(), 1);
  EXPECT_FALSE(window.handler().last_pushed.has_value());
  EXPECT_FALSE(window.handler().last_popped.has_value());

  window.Push(0, 1);
  window.Push(0, 2);
  window.Push(0, 3);
  window.Push(0, 4);
  EXPECT_EQ(window.size(), 4);
  EXPECT_FALSE(window.empty());
  EXPECT_EQ(window.num_popped(), 0);
  EXPECT_EQ(window.handler().last_pushed.value().first, 0);
  EXPECT_EQ(window.handler().last_pushed.value().second, 4);
  EXPECT_FALSE(window.handler().last_popped.has_value());

  window.Push(5, 5);
  EXPECT_EQ(window.size(), 1);
  EXPECT_EQ(window.num_popped(), 4);
  EXPECT_EQ(window.handler().last_pushed.value().first, 5);
  EXPECT_EQ(window.handler().last_pushed.value().second, 5);
  EXPECT_EQ(window.handler().last_popped.value().first, 0);
  EXPECT_EQ(window.handler().last_popped.value().second, 4);

  window.PushTime(5);
  EXPECT_EQ(window.size(), 1);
  EXPECT_EQ(window.num_popped(), 4);
  EXPECT_EQ(window.handler().last_pushed.value().first, 5);
  EXPECT_EQ(window.handler().last_pushed.value().second, 5);
  EXPECT_EQ(window.handler().last_popped.value().first, 0);
  EXPECT_EQ(window.handler().last_popped.value().second, 4);

  window.PushTime(6);
  EXPECT_EQ(window.size(), 0);
  EXPECT_EQ(window.num_popped(), 5);
  EXPECT_EQ(window.handler().last_pushed.value().first, 5);
  EXPECT_EQ(window.handler().last_pushed.value().second, 5);
  EXPECT_EQ(window.handler().last_popped.value().first, 5);
  EXPECT_EQ(window.handler().last_popped.value().second, 5);
}

TEST(TimeMovingWindowTest, Xarray) {
  SimpleTimeMovingWindow<xt::xarray<int>, true> window;
  window.ResizeWindow(1);
  EXPECT_EQ(window.window_size(), 1);
  EXPECT_FALSE(window.handler().last_pushed.has_value());
  EXPECT_FALSE(window.handler().last_popped.has_value());

  window.Push(0, {{1, 2}, {3, 4}});
  window.Push(0, {2});
  window.Push(0, {3});
  window.Push(0, {{1, 0}, {0, 1}});
  EXPECT_EQ(window.size(), 4);
  EXPECT_FALSE(window.empty());
  EXPECT_EQ(window.num_popped(), 0);
  EXPECT_EQ(window.handler().last_pushed.value().first, 0);
  EXPECT_EQ(window.handler().last_pushed.value().second, xt::xarray<int>({{1, 0}, {0, 1}}));
  EXPECT_FALSE(window.handler().last_popped.has_value());

  window.Push(5, {5});
  EXPECT_EQ(window.size(), 1);
  EXPECT_EQ(window.num_popped(), 4);
  EXPECT_EQ(window.handler().last_pushed.value().first, 5);
  EXPECT_EQ(window.handler().last_pushed.value().second, xt::xarray<int>({5}));
  EXPECT_EQ(window.handler().last_popped.value().first, 0);
  EXPECT_EQ(window.handler().last_popped.value().second, xt::xarray<int>({{1, 0}, {0, 1}}));

  window.PushTime(5);
  EXPECT_EQ(window.size(), 1);
  EXPECT_EQ(window.num_popped(), 4);
  EXPECT_EQ(window.handler().last_pushed.value().first, 5);
  EXPECT_EQ(window.handler().last_pushed.value().second, xt::xarray<int>({5}));
  EXPECT_EQ(window.handler().last_popped.value().first, 0);
  EXPECT_EQ(window.handler().last_popped.value().second, xt::xarray<int>({{1, 0}, {0, 1}}));

  window.PushTime(6);
  EXPECT_EQ(window.size(), 0);
  EXPECT_EQ(window.num_popped(), 5);
  EXPECT_EQ(window.handler().last_pushed.value().first, 5);
  EXPECT_EQ(window.handler().last_pushed.value().second, xt::xarray<int>({5}));
  EXPECT_EQ(window.handler().last_popped.value().first, 5);
  EXPECT_EQ(window.handler().last_popped.value().second, xt::xarray<int>({5}));
}

TEST(TimeMovingWindowMaxTest, Basic) {
  std::array<int, 11> arr {8, 5, 10, 7, 9, 4, 3, 15, 12, 90, 13};
  TimeMovingWindowMax<int> window;
  window.ResizeWindow(4);

  std::vector<int> res;
  for (int i = 0; i < arr.size(); ++i) {
    window.Push(i, arr[i]);
    ASSERT_TRUE(window.has_value());
    res.push_back(window.value());
  }

  EXPECT_EQ(res, std::vector({8, 8, 10, 10, 10, 10, 9, 15, 15, 90, 90}));
}

TEST(TimeMovingWindowMinTest, Basic) {
  std::array<int, 11> arr {8, 5, 10, 7, 9, 4, 3, 15, 12, 90, 13};
  TimeMovingWindowMin<int> window;
  window.ResizeWindow(4);

  std::vector<int> res;
  for (int i = 0; i < arr.size(); ++i) {
    window.Push(i, arr[i]);
    ASSERT_TRUE(window.has_value());
    res.push_back(window.value());
  }

  EXPECT_EQ(res, std::vector({8, 5, 5, 5, 5, 4, 3, 3, 3, 3, 12}));
}

TEST(TimeMovingWindowMaxTest, PushTime) {
  TimeMovingWindowMax<int> window;
  window.ResizeWindow(2);
  EXPECT_FALSE(window.has_value());

  window.Push(0, 0);
  ASSERT_TRUE(window.has_value());
  EXPECT_EQ(window.value(), 0);

  window.Push(1, 1);
  ASSERT_TRUE(window.has_value());
  EXPECT_EQ(window.value(), 1);

  window.Push(2, 2);
  ASSERT_TRUE(window.has_value());
  EXPECT_EQ(window.value(), 2);

  window.Push(3, 1);
  ASSERT_TRUE(window.has_value());
  EXPECT_EQ(window.value(), 2);

  window.Push(4, 0);
  ASSERT_TRUE(window.has_value());
  EXPECT_EQ(window.value(), 1);

  window.PushTime(5);
  ASSERT_TRUE(window.has_value());
  EXPECT_EQ(window.value(), 0);

  window.PushTime(6);
  EXPECT_FALSE(window.has_value());
}

TEST(TimeMovingWindowMaxTest, String) {
  std::array<std::string, 11> arr {"8", "5", "a", "7", "9", "4", "3", "d", "b", "e", "c"};
  TimeMovingWindowMax<std::string> window;
  window.ResizeWindow(4);

  std::vector<std::string> res;
  for (int i = 0; i < arr.size(); ++i) {
    window.Push(i, arr[i]);
    ASSERT_TRUE(window.has_value());
    res.push_back(window.value());
  }

  std::vector<std::string> ans {"8", "8", "a", "a", "a", "a", "9", "d", "d", "e", "e"};
  EXPECT_EQ(res, ans);
}

}  // namespace coin2::base::math
